"""
    Title

    author: wxz
    date: 2021-12-28
    github: https://github.com/xinzwang

    ** SoftWare ISP Hyperparameter
        --  version: 1.0.0

              type: numpy.ndarray()
              shape: 20

              black level correction   params: 2  [black_level, white_level]
              demosaik                 params: 0
              white balance            params: 3  [gain_1, gain_2, gain_3]
              color correction         params: 9  [c_00, c_01, c_02, c_10, c_11, c_12, c_31, c_32,c_33]
              gamma correction         params: 2  [mul, gamma]
              tone curve               params: 4  [mul, kernel_w, kernel_h, sigma]

"""
import os
import glob
import re
import json
import numpy as np

from isp.blocks.black_level_correction import black_level_correction
from isp.blocks.color_correction import color_correction
from isp.blocks.demosaik import demosaik
from isp.blocks.gamma_correction import gamma_correction
from isp.blocks.tone_curve import tone_curve
from isp.blocks.white_balance import channel_gain_white_balance


class ISP(object):
    def __init__(self):
        pass

    @staticmethod
    def init_params():
        # init params of isp
        default_params = np.array(
            [512, 16383, 1.9296875, 1.0, 2.26171875, .9020, -.2890, -.0715, -.4535, 1.2436, .2348, -.0934, .1919, .7086,
             80.0, 2.2, 0.1, 3.0, 3.0, 1.0])
        return default_params

    def forward(self, raw, p):
        # version 1.0.0     p->params
        assert raw is not None, print('[ERROR] input raw data is None')
        assert p.shape == (20,), print('[ERROR] params shape error. shape:', p.shape)
        img = black_level_correction(raw, p[0], p[1])
        img = demosaik(img)
        img = channel_gain_white_balance(img, p[2:5])
        img = color_correction(img, p[5:14].reshape(3, 3))
        img = gamma_correction(img, p[14], p[15])
        img = tone_curve(img, p[16], p[17:19], p[19])
        return img

    @staticmethod
    def save_params(p=None, path="./runs/isp", suffix=""):
        # save isp params to json file
        a = {
            "black_level_correction": {
                "black_level": p[0],
                "white_level": p[1]
            },
            "white_balance": p[2:5].tolist(),
            "color_correction": p[5:14].tolist(),
            "gamma_correction": {
                "mul": p[14],
                "gamma": p[15]
            },
            "tone_curve": {
                "mul": p[16],
                "kernel": p[17:19].tolist(),
                "sigma": p[19]
            }
        }
        path = path.replace('/', os.sep)
        if not path.split('.')[-1] == 'json':
            files = glob.glob(str(path + "\\*.json"))
            num = [int(re.search('[0-9]+', x.split('.')[-2]).group()) for x in files]  # find suffix num
            num = 0 if len(num) == 0 else max(num) + 1
            path = path + os.sep + "params_" + suffix + str(num) + ".json"
        with open(path, 'w') as f:
            f.write(json.dumps(a))
        return

    @staticmethod
    def load_params(path):
        # load isp params from json file
        a = None
        with open(path) as f:
            a = json.loads(f.read())
        p = [
            a["black_level_correction"]['black_level'],
            a["black_level_correction"]['white_level'],
            a["white_balance"],
            a["color_correction"],
            a["gamma_correction"]["mul"],
            a["gamma_correction"]["gamma"],
            a["tone_curve"]["mul"],
            a["tone_curve"]["kernel"],
            a["tone_curve"]["sigma"]
        ]
        p = flatten(p)
        p = np.array(p)
        return p


# ---------------- utils function ----------------
# Flatten the list to 1-D
flatten = lambda x: [y for l in x for y in flatten(l)] if type(x) is list else [x]

if __name__ == '__main__':
    a = ISP()
    p = a.init_params()



